'''
pyrwi.Fl_Plot - A new pyFLTK widget that use matplotlib to make nice GUI plots
            This could be filed under pyrwi.gui but (1) it is pretty large
            and this helps keep the pyrwi.gui file to a manageable length
            and (2) this could be a useful tool for others who don't need
            the whole pyrwi set, so trying to keep it separate.
'''
import fltk

try:
    import matplotlib as mpl
    mpl.use('FltkAgg')
    #Agg seems to be about 10us faster than FltkAgg
    #But FltkAgg will allow us to still make plots from text only windows
except:
    print "Error importing FltkAgg. Fl_Plot may not work."

import pylab as pl

from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.backend_bases import FigureCanvasBase
from matplotlib.backends.backend_fltkagg import special_key

class Fl_Plot(fltk.Fl_Widget, FigureCanvasAgg):
    def __init__(self,x,y,w,h,label, use_test_img=False):
        '''
        Fl_Plot constructor
        @param  x : X Location of plot
        @param  y : Y Location of plot
        @param  w : Width of plot
        @param  h : Height of plot
        @param  label : Widget label
        @param  user_test_img : Replace plot with test image. For debugging only
        '''
        #--- Configuration Options
        #How much space should be given between maximal data points
        #and edge of plot. The numbers must be >= 0
        self.x_padding = 0.05 
        self.y_padding = 0.2 
        
        #--- Setup the FLTK Widget
        fltk.Fl_Widget.__init__(self, x, y, w, h, label)
        self.x = x
        self.y = y
        
        #--- Create an MPL figure
        self.figure = pl.Figure()
        #add_subplot() is the slowest part of the whole __init__
        self.ax = self.figure.add_subplot(111)
        self.lines = [] #List of all lines on this plot
        
        #Figure Canvas
        FigureCanvasAgg.__init__(self,self.ax.figure)
        
        #--- Internal data
        self._oldsize=(None,None) #Remember last widget size
        self._draw_overlay = False #For handle() to remember state
        self._button = None #For handle() to remember state
        self._key = None #For handle() to remember state

        self.background = None #Cache plot background
        #Every time the canvas is re-drawn, cache the new background
        self.ax.figure.canvas.mpl_connect('draw_event', self._update_background)
        
        #self.get_renderer() is quite slow, so we will cache the result
        self.renderer = self.get_renderer() #For figure.draw(self.renderer)
                
        #--- Make a test image the same size as the plot
        self.test_img=use_test_img
        if self.test_img :
            self.make_test_data()

    def make_test_data(self) :
        '''
        Generate an RGBA image that can replace the plot
        '''
        w = self.w()
        h = self.h()
        sz = w * h * 4 #4 - this is Red, Green, Blue, and an Alpha
        self.rgba = [128]*sz #Init image
        i = 0
        for y in xrange(h) :
            for x in xrange(w) :
                self.rgba[i+0] = 95+y #Red
                self.rgba[i+1] = 128 #Green
                self.rgba[i+2] = 128 #Blue
                self.rgba[i+3] = 0 #Alpha
                i += 4

    def set_ylim(self, limits):
        '''
        Set the limits for the yaxis; v = [ymin, ymax]::

          set_ylim((valmin, valmax))
          set_ylim(valmin, valmax)
          set_ylim(ymin=1) # ymax unchanged
          set_ylim(ymax=1) # ymin unchanged

        Keyword arguments:

          *ymin*: scalar
            the min of the ylim
          *ymax*: scalar
            the max of the ylim
          *emit*: [ True | False ]
            notify observers of lim change

        Returns the current ylimits as a length 2 tuple

        ACCEPTS: len(2) sequence of floats
        '''
        #TODO: allow user to fix axis
        pass
    def set_xlim(self, limits):
        #TODO: allow user to fix axis
        pass

    
    def legend(self, *args, **kwargs):
        self.ax.legend(*args, **kwargs)
        
        self.draw(force=True)
        for l in self.lines : self.ax.draw_artist(l)
        self.redraw()
        
    def clear(self) :
        '''
        Clear the plot and the stored lines.
        '''
        self.ax.clear()
        self.lines = [] #No lines are on the plot now
        self.draw(force=True) #Force re-draw of axis
        
    def plot(self, *args, **kwargs) :
        '''
        Call pylab.plot for this figure. Also adds each line object to self.lines
        @@param *args, **kwargs: Same as pylab.axes.plot()
        @@return: Same as pylab.axes.plot()
        '''
        lines = self.ax.plot(animated=True, *args, **kwargs)

        #MPL has computed axis limits that fit the data tightly
        #resize the axis according to the padding configured here
        self._setAxisLimits( self.ax.get_xlim() + self.ax.get_ylim() )
        
        self.draw(force=True) #Force re-draw of axis
        
        #Add new lines to set of old lines, for use in update()
        self.lines += lines
        #Redraw all lines
        for l in self.lines :
            self.ax.draw_artist(l) #Add line to plot
        
        self.redraw() #Put the lines on the plot
        return lines
    
    def update(self, new_x, *new_y) :
        '''
        Update all the lines (using the same X-vector)
        @param new_x : New X vector. Used for all Y vectors
        @param *new_y : Y-Vectors for each line on the plot
        '''
        assert len(new_y) == len(self.lines), "Y data set must be same size as number of lines"

        if self.background is None:
            # Cache a clean background - This should never be called
            # because the background is not clean here. Should have been
            # cached already on the last canvas draw_event
            self.background = self.copy_from_bbox(self.ax.bbox)

        #--- Adjust Axis if necissary
        #Performance notes:
        # - min/max could be done in for (l,y) but then we have to re-call self.update
        # - min/max could be done here with one loop, but this is faster
        min_y = min( min(new_y)) #Inner min returns the list with the min in it
        max_y = max( max(new_y)) #outer max picks the max element from the max list
        axis_new = (min(new_x), max(new_x), min_y, max_y)
        newAxis = False
        if self._checkAxisLimits( *axis_new ) :
            newAxis = True
            #See: http://matplotlib.sourceforge.net/examples/animation/strip_chart_demo.html
            self._setAxisLimits(axis_new)
            self.ax.figure.canvas.draw()
        
        #Restore a blank set of axis to the figure (doesn't plot anyting)
        self.restore_region( self.background )
        if newAxis :
            #Must draw full plot so axis are drawn too
            #This only draws the BG, so must come before
            #re-drawing the lines
            self.draw(force=True)
        
        # update the data
        for (l,y) in zip(self.lines, new_y) :
            assert len(new_x) == len(y), "All y vectors must be same length as x vector"
            l.set_data([new_x,y]) #this is quite a bit faster than set_ydata
            self.ax.draw_artist(l) #Re-draw just this line

        #Still need to blit or I am left with a blank plot sometimes
        self.blit() #Splatter only the plot to the screen
    
    def _checkAxisLimits(self, xmin, xmax, ymin, ymax) :
        '''
        Determine if a new axis need to be re-drawn.
        The function call overhead for this is minimal
        '''
        if xmin < self.axis_limits[0] : return True
        elif xmax > self.axis_limits[1] : return True
        elif ymin < self.axis_limits[2] : return True
        elif ymax > self.axis_limits[3] : return True
        return False
        
    def _setAxisLimits(self, axis_new) :
        xmargin = self.x_padding*(axis_new[1] - axis_new[0])
        ymargin = self.y_padding*(axis_new[3] - axis_new[2])
        
        self.ax.set_xlim( axis_new[0]-xmargin, axis_new[1]+xmargin )
        self.ax.set_ylim( axis_new[2]-ymargin, axis_new[3]+ymargin )

        #Cache the current axis limits, so we can tell if they change
        xaxis_old = self.ax.get_xlim()
        yaxis_old = self.ax.get_ylim()
        self.axis_limits = xaxis_old + yaxis_old

    def _update_background(self, event) :
        '''
        Store a background with no artists on it. 
        This is called from the canvas draw event by MPL.
        '''
        self.background = self.copy_from_bbox(self.ax.bbox)

    def draw(self, force=False):
        '''
        Draw the plot to the screen.
        @param force: OPTIONAL If true, re-draw the widget even if it wasn't re-sized
        '''

        newsize=(self.w(),self.h())
        if(self._oldsize != newsize or force):
            self._oldsize =newsize
            
            #Set figure size: compute desired figure size in inches
            dpival = self.figure.dpi
            winch = self.w()/dpival
            hinch = self.h()/dpival
            self.figure.set_size_inches(winch,hinch)        
            
            #Re-generate figure image
            FigureCanvasAgg.draw(self)
            #TODO: Problem here!
        
        #--- Draw figure into widget
        t1,t2,w,h = self.figure.bbox.bounds
        
        fltk.fl_push_no_clip() #Don't clip anything, draw anywhere
        if True or not self.test_img :
            #Take figure image and put it on the screen
            fltk.fl_draw_image(self.buffer_rgba(0,0),self.x,self.y,int(w),int(h),4,0)
        else :
            #Write test image
            fltk.fl_draw_image(self.rgba,self.x,self.y,int(w),int(h),4,0)
            for i in xrange(len(self.rgba)) :
                self.rgba[i] += 10
                if self.rgba[i] > 255 : self.rgba[i] = 0

        fltk.fl_pop_clip() #Restore orig clip region        
        
    def blit(self):
        '''
        Draw only the bare minimum to the screen
        '''

        #Location of axis within figure
        x = int(self.ax.bbox.x0)
        y = int(self.ax.bbox.y0)

        t1,t2,w,h = self.ax.figure.bbox.bounds
        wo= int(w) #Width of actual RGB image
        w = int(w)-x #Width I want to plot
        h = int(h)-y #Height I want to plot

        fltk.fl_push_no_clip() #Don't clip anything, draw anywhere
        
        if self.test_img :
            #Splatter a test image up first
            fltk.fl_draw_image(self.rgba,self.x+x,self.y+y,w,h,4,wo*4)
            for i in xrange(len(self.rgba)) :
                self.rgba[i] += 10
                if self.rgba[i] > 255 : self.rgba[i] = 0

        #Write only axis area to screen
        fltk.fl_draw_image(self.buffer_rgba(x,y),self.x+x,self.y+y,w,h,4,wo*4)
        #fl_draw_image starts at x,y in the buffer.
        # Then it counts w*4 elements over outputing each pixel
        # Next it jumps from the starting element to wo*4 elements later and starts a new row
        # repeat until h rows have been covered
        
        #Write entire image to screen
        #fltk.fl_draw_image(self.buffer_rgba(0,0),self.x,self.y,wo,h+y,4,0)
        
        
        fltk.fl_pop_clip() #Restore clipping to the way I found it

    #TODO: May need to add these for the sake of FigureCanvasAgg
    #def resize(self,size):
        #w, h = size
        ## compute desired figure size in inches
        #dpival = self.figure.dpi
        #winch = w/dpival
        #hinch = h/dpival
        #self.figure.set_size_inches(winch,hinch)

    def handle(self, event):
        '''
        Event handler
        @param event: Integer indicating what event was raised (http://pyfltk.sourceforge.net/docs/CH6_Events.html#events)
        '''
        #TODO: Does any of this work with the widget?
        x=fltk.Fl.event_x()
        y=fltk.Fl.event_y()
        yf=self.figure.bbox.height - y
        if event == fltk.FL_FOCUS or event == fltk.FL_UNFOCUS:
            return 1
        elif event == fltk.FL_KEYDOWN:
            ikey= fltk.Fl.event_key()
            if(ikey<=255):
                self._key=chr(ikey)
            else:
                try:
                    self._key=special_key[ikey]
                except:
                    self._key=None
            FigureCanvasBase.key_press_event(self, self._key)
            return 1
        elif event == fltk.FL_KEYUP:
            FigureCanvasBase.key_release_event(self, self._key)
            self._key=None
        elif event == fltk.FL_PUSH:
            self.window().make_current()
            if fltk.Fl.event_button1():
                self._button = 1
            elif  fltk.Fl.event_button2():
                self._button = 2
            elif  fltk.Fl.event_button3():
                self._button = 3
            else:
                self._button = None

            if self._draw_overlay:
                self._oldx=x
                self._oldy=y
            if fltk.Fl.event_clicks():
                FigureCanvasBase.button_press_event(self, x, yf, self._button)
                return 1
            else:
                FigureCanvasBase.button_press_event(self, x, yf, self._button)
                return 1
        elif event == fltk.FL_ENTER:
            self.take_focus()
            return 1
        elif event == fltk.FL_LEAVE:
            return 1
        elif event == fltk.FL_MOVE:
            FigureCanvasBase.motion_notify_event(self, x, yf)
            return 1
        elif event == fltk.FL_DRAG:
            self.window().make_current()
            if self._draw_overlay:
                self._dx=fltk.Fl.event_x()-self._oldx
                self._dy=fltk.Fl.event_y()-self._oldy
                fltk.fl_overlay_rect(self._oldx,self._oldy,self._dx,self._dy)
            FigureCanvasBase.motion_notify_event(self, x, yf)
            return 1
        elif event == fltk.FL_RELEASE:
            self.window().make_current()
            if self._draw_overlay:
                fltk.fl_overlay_clear()
            FigureCanvasBase.button_release_event(self, x, yf, self._button)
            self._button = None
            return 1
        return 0

    def start_event_loop(self,timeout):
        FigureCanvasBase.start_event_loop_default(self,timeout)
    start_event_loop.__doc__=FigureCanvasBase.start_event_loop_default.__doc__

    def stop_event_loop(self):
        FigureCanvasBase.stop_event_loop_default(self)
    stop_event_loop.__doc__=FigureCanvasBase.stop_event_loop_default.__doc__

